#!/usr/bin/env python
"""Generate code for wandb SDK.

Usage:
    ./tools/generate-tool.py --generate
    ./tools/generate-tool.py --check
    ./tools/generate-tool.py --generate --check
    ./tools/generate-tool.py --generate --check wandb/sdk/lib/_wburls_generated.py
"""

import argparse
import contextlib
import filecmp
import os
import subprocess
import tempfile
from pathlib import Path, PurePath
from typing import Iterator, List

parser = argparse.ArgumentParser()
parser.add_argument("--generate", action="store_true", help="generate files")
parser.add_argument(
    "--no-format",
    dest="format",
    action="store_false",
    default=True,
    help="format files",
)
parser.add_argument(
    "--check", action="store_true", help="check if generated files are up-to-date"
)
parser.add_argument("files", nargs="*")
args = parser.parse_args()


GENERATE_SUFFIX = "_generate.py"
GENERATED_SUFFIX = "_generated.py"


def get_paths() -> List[PurePath]:
    paths: List[PurePath] = []
    if not args.files:
        exclude_dirs = {"vendor", "__pycache__"}
        root_dir = Path(__file__).resolve().parent.parent / "wandb"
        for base, subdirs, files in os.walk(root_dir):
            # Don't walk into excluded subdirectories
            subdirs[:] = list(set(subdirs) - exclude_dirs)
            for fname in files:
                if fname.endswith(GENERATE_SUFFIX):
                    paths.append(PurePath(base, fname))
    for f in args.files:
        paths.append(Path(f))
    return paths


def generate_file(generate_path: Path, output_path: Path) -> None:
    status, output = subprocess.getstatusoutput(f"python {generate_path}")
    assert status == 0, f"Error: {output}"
    with open(output_path, "w") as f:
        f.write("# DO NOT EDIT -- GENERATED BY: `generate-tool.py --generate`")
        f.write(output)


def generate_files(paths: List[PurePath]) -> None:
    for p in paths:
        output_path = p.parent / str(p).replace(GENERATE_SUFFIX, GENERATED_SUFFIX)
        print(f"INFO: Generating {output_path}...")
        generate_file(p, output_path)


def format_file(filename: Path) -> None:
    status, output = subprocess.getstatusoutput(f"black {filename}")
    assert status == 0, f"Error: {output}"


def format_files(paths: List[PurePath]) -> None:
    for p in paths:
        output_path = p.parent / str(p).replace(GENERATE_SUFFIX, GENERATED_SUFFIX)
        print(f"INFO: Formatting {output_path}...")
        format_file(output_path)


@contextlib.contextmanager
def temp_fname() -> Iterator[PurePath]:
    try:
        f = tempfile.NamedTemporaryFile(delete=False)
        tmp_name = f.name
        f.close()
        yield Path(tmp_name)
    finally:
        os.unlink(tmp_name)


def check_files(paths: List[PurePath]) -> None:
    for p in paths:
        generated_path = p.parent / str(p).replace(GENERATE_SUFFIX, GENERATED_SUFFIX)
        print(f"INFO: Checking {generated_path}...")
        with temp_fname() as temp_file:
            generate_file(p, temp_file)
            format_file(temp_file)
            assert filecmp.cmp(
                generated_path, temp_file
            ), f"expected: {open(temp_file).read()}"


def main() -> None:
    path_list = get_paths()

    if args.generate:
        generate_files(path_list)
        if args.format:
            format_files(path_list)

    if args.check:
        check_files(path_list)


if __name__ == "__main__":
    main()
